import random

def arrange_array_by_pivot(array, pivot):
    """ 将数组分为3部分， 第1部分元素小于pivot, 第二部分元素等于pivot，第三部分元素大于pivot """
    i = 0
    j = len(array) - 1

    while i < j:
        if array[i] < pivot:
            i += 1
        elif array[i] >= pivot:
            temp = array[j]
            array[j] = array[i]
            array[i] = temp
            j -= 1
        
    if array[i] < pivot:
        i += 1
    
    k = i
    S1 = array[0:i]
    j = len(array) - 1

    while i < j:
        if array[i] == pivot:
            i += 1
        else:
            temp = array[j]
            array[j] = array[i]
            array[i] = temp
            j -= 1
    if array[i] == pivot:
        i += 1
    
    S2 = array[k:i]
    S3 = array[i:]

    return S1, S2, S3
    

def selection(array, k):
    """ 查找数组中第k小的元素 """
    if len(array) <= 1:
        return array[0]
    if k > len(array):
        raise Exception('Out of range!')
    
    pivot = array[0]
    is_good_pivot = False

    while is_good_pivot != True:
        pivot = random.choice(array)

        S1, S2, S3 = arrange_array_by_pivot(array, pivot)
        if len(S1) <= int((3/4)*len(array)) and len(S3) <= int((3/4)*len(array)):
            is_good_pivot = True
    
    if len(S1) > 0 and k <= len(S1):
        return selection(S1, k)
    elif len(S2) > 0 and k <= len(S1) + len(S2):
        return pivot
    else:
        return selection(S3, k - len(S1) - len(S2))
    raise Exception('error')

array = [138, 132, 119, 149, 116, 124, 121, 144, 203, 105, 180, 185, 165,
         113, 173, 117, 176, 196, 122, 142, 172, 162, 190, 118, 150, 128, 186, 177, 189, 183,
         189, 123, 190, 133, 147, 182, 170, 183, 157, 124, 194, 115, 203, 203, 200, 152, 178,
         186, 118, 183, 202]

k = int((len(array) + 1) / 2)
print(f"median number is :{k}")
median = selection(array, k)
print('The element selected is', median)
